import numpy as np


def model(X, W):
    Z = np.dot(X, W)
    return Z


def sigmoid(Z):
    A = 1. / (1. + np.exp(- Z))
    return A


def hypo(X, W):
    Z = model(X, W)
    A = sigmoid(Z)
    return A


def error(H, Y):
    E = H - Y
    return E


def cost(A, Y):
    m = Y.shape[0]
    ce = np.dot(Y.T, np.log(A)) + np.dot(1. - Y.T, np.log(1. - A))
    j = ce[0, 0] / - m
    return j


def delta_W(X, E):
    m = X.shape[0]
    dW = np.dot(X.T, E) / m
    return dW


def update_W(W, dW, alpha):
    W -= alpha * dW


if '__main__' == __name__:
    from python_ai.category.data.data4logistic_regr import X, Y
    from mpl_toolkits import mplot3d
    import matplotlib.pyplot as plt

    np.random.seed(1)

    X_mean = X.mean(axis=0)
    X_sigma = X.std(axis=0)
    X -= X_mean  # (M, 3); (3,) => (M, 3)
    X /= X_sigma
    X = np.concatenate([
        np.ones([X.shape[0], 1]),
        X
    ], axis=1)

    W = np.random.normal(0., 1., [X.shape[1], 1])

    ITERS = 1000
    ALPHA = 1e-2
    j_his = []
    for i in range(ITERS):
        A = hypo(X, W)
        E = error(A, Y)
        j = cost(A, Y)
        j_his.append(j)
        dW = delta_W(X, E)
        update_W(W, dW, ALPHA)

    print('W:', W)

    plt.figure(figsize=[12, 6])
    plt_params = dict(spr=1, spc=2, spn=0)
    plt_params['spn'] += 1
    ax = plt.subplot(plt_params['spr'], plt_params['spc'], plt_params['spn'], projection='3d')
    x1_v = X[:, 1]
    x2_v = X[:, 2]
    x3_v = X[:, 3]
    y_v = Y[:, 0]
    ax.scatter3D(x1_v[y_v > 0.5], x2_v[y_v > 0.5], x3_v[y_v > 0.5], color='red')
    ax.scatter3D(x1_v[y_v < 0.5], x2_v[y_v < 0.5], x3_v[y_v < 0.5], color='blue')

    x1_samples = np.linspace(X[:, 1].min(), X[:, 1].max(), 5)
    x2_samples = np.linspace(X[:, 2].min(), X[:, 2].max(), 5)
    xx, yy = np.meshgrid(x1_samples, x2_samples)
    xxr = xx.ravel()
    yyr = yy.ravel()
    Samples = np.concatenate([
        np.ones_like(xxr).reshape(-1, 1),
        xxr.reshape(-1, 1),
        yyr.reshape(-1, 1),
    ], axis=1)
    zz = (np.dot(Samples, W[:-1, :]) / - W[-1, 0]).reshape(xx.shape)
    print('xx yy zz', xx.shape, yy.shape, zz.shape)
    ax.plot_surface(xx, yy, zz)

    plt_params['spn'] += 1
    plt.subplot(plt_params['spr'], plt_params['spc'], plt_params['spn'])
    plt.title('Cost history')
    plt.plot(j_his)
    plt.show()
